from flask import Flask, request, render_template
import ollama

app = Flask(__name__)

@app.route('/', methods=['GET', 'POST'])
def chat():
    if request.method == 'POST':
        prompt = request.form['prompt']
        response = ollama.chat(
            model='llama3',
            messages=[{'role': 'user', 'content': prompt}]
        )
        return render_template('index.html', response=response['message']['content'])
    return render_template('index.html')

if __name__ == '__main__':
    app.run(port=5000, debug=True)